import numpy as np
from typing import List, Dict, Tuple

class AgentPruning:
    """Implements pruning methods for agent's cognitive decision making"""
    
    def __init__(self):
        self.pruning_threshold = 0.6  # Configurable threshold
        self.memory_limit = 100       # Max nodes to keep
        
    def prune_perception_nodes(self, nodes: List[Dict], features: List[str]) -> List[Dict]:
        """Prune perception nodes based on importance and recency"""
        if len(nodes) <= self.memory_limit:
            return nodes
            
        scores = []
        for node in nodes:
            # Calculate importance score
            importance = node.get('importance_score', 0)
            recency = self._calculate_recency(node['timestamp'])
            relevance = self._feature_relevance(node['features'], features)
            
            score = 0.4 * importance + 0.3 * recency + 0.3 * relevance
            scores.append(score)
            
        # Keep only highest scoring nodes
        indices = np.argsort(scores)[-self.memory_limit:]
        return [nodes[i] for i in indices]
    
    def prune_planning_options(self, plans: List[Dict], 
                             current_state: Dict) -> List[Dict]:
        """Prune planning options based on feasibility and potential value"""
        if not plans:
            return []
            
        viable_plans = []
        for plan in plans:
            # Calculate viability score
            energy_viable = plan['energy_required'] <= current_state['energy']
            knowledge_viable = current_state['knowledge'] >= plan.get('knowledge_required', 0)
            success_prob = plan.get('success_probability', 0)
            
            score = (0.4 * energy_viable + 
                    0.3 * knowledge_viable + 
                    0.3 * success_prob)
                    
            if score > self.pruning_threshold:
                viable_plans.append(plan)
                
        return viable_plans
    
    def prune_reasoning_paths(self, paths: List[Dict], 
                            constraints: Dict) -> List[Dict]:
        """Prune reasoning paths based on constraints and expected outcomes"""
        valid_paths = []
        
        for path in paths:
            # Validate against constraints
            meets_energy = path['energy_cost'] <= constraints['max_energy']
            meets_time = path.get('estimated_time', 0) <= constraints['max_time']
            meets_safety = path.get('safety_score', 1) >= constraints['min_safety']
            
            # Calculate expected value
            expected_value = (path['success_probability'] * 
                            path.get('reward', 1) -
                            path['energy_cost'])
            
            if meets_energy and meets_time and meets_safety:
                if expected_value > self.pruning_threshold:
                    valid_paths.append(path)
        
        return valid_paths
    
    def _calculate_recency(self, timestamp: str) -> float:
        """Calculate recency score based on timestamp"""
        # Implementation based on time difference
        return 0.8  # Placeholder
        
    def _feature_relevance(self, features: Dict, 
                          relevant_features: List[str]) -> float:
        """Calculate feature relevance score"""
        if not relevant_features:
            return 1.0
            
        matches = sum(1 for f in relevant_features if f in features)
        return matches / len(relevant_features) 